package funcdrawer

import (
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	"image/png"
	"io"
)

type Function interface {
	Eval(x float64) float64
}

const (
	thicknessFactor = 250
	xresolution     = 0.01
)

// Why can't this be a constant?
var (
	background    = color.RGBA{240, 240, 240, 0}
	axesColor     = color.Black
	functionColor = color.RGBA{255, 0, 0, 255}
)

// Draws the background of the image.
func drawBackground(m *image.RGBA) {
	draw.Draw(m, m.Bounds(), image.NewUniform(background), image.ZP, draw.Src)
}

// Draws the x and y axes on the image m given the viewport.
func drawAxes(m *image.RGBA, viewport image.Rectangle) {
	imgBounds := m.Bounds()

	computeZeroPoint := func(minVal, maxVal, minBound, maxBound int) int {
		relPos := float64(-minVal) / float64(maxVal-minVal)
		return int(relPos * float64(maxBound-minBound))
	}

	if viewport.Min.X < 0 && viewport.Max.X > 0 {
		xAxPos := computeZeroPoint(
			viewport.Min.X, viewport.Max.X,
			m.Bounds().Min.X, m.Bounds().Max.X)
		thickness := m.Bounds().Dx() / thicknessFactor
		ax := image.Rect(
			xAxPos-thickness/2, imgBounds.Min.Y,
			xAxPos+thickness/2+1, imgBounds.Max.Y)
		draw.Draw(m, ax, image.NewUniform(axesColor), image.ZP, draw.Src)
	}

	if viewport.Min.Y < 0 && viewport.Max.Y > 0 {
		yAxPos := m.Bounds().Max.Y - computeZeroPoint(
			viewport.Min.Y, viewport.Max.Y,
			m.Bounds().Min.Y, m.Bounds().Max.Y)
		thickness := m.Bounds().Dy() / thicknessFactor
		ax := image.Rect(
			imgBounds.Min.X, yAxPos-thickness/2,
			imgBounds.Max.X, yAxPos+thickness/2+1)
		draw.Draw(m, ax, image.NewUniform(axesColor), image.ZP, draw.Src)
	}
}

// A 2-dimension point in the integers.
type Point2D struct {
	x, y int
}

// Draws a straight line between two given points in the given image.
func drawLine(m *image.RGBA, a, b *Point2D) {
	order := func(a, b *Point2D) (*Point2D, *Point2D) {
		if a.x < b.x {
			return a, b
		}
		return b, a
	}
	minX, maxX := order(a, b)
	d := xresolution * float64(maxX.y-minX.y) / float64(maxX.x-minX.x)
	for x, y := float64(minX.x), float64(minX.y); x <= float64(maxX.x); x, y = x+xresolution, y+d {
		m.Set(int(x), int(y), functionColor)
	}
}

// Draws all the values in the given viewport for the passed function on the image m.
func drawValues(m *image.RGBA, viewport image.Rectangle, f Function) {
	bounds := m.Bounds()

	x0, xm := bounds.Min.X, bounds.Max.X
	y0, ym := bounds.Min.Y, bounds.Max.Y
	vx0, vxm := float64(viewport.Min.X), float64(viewport.Max.X)
	vy0, vym := float64(viewport.Min.Y), float64(viewport.Max.Y)

	var lastPoint *Point2D = nil
	for x := x0; x < xm; x++ {
		vx := float64(x-x0)*(vxm-vx0)/float64(xm-x0) + vx0
		vy := f.Eval(vx)
		y := int((vy-vy0)*float64(y0-ym)/(vym-vy0) + float64(ym))
		newPoint := &Point2D{x, y}
		if lastPoint != nil {
			drawLine(m, lastPoint, newPoint)
		}
		lastPoint = newPoint
	}
}

// Draws an image containing the values in the given viewport of f.
// The it encodes the image either to jpeg or png depending on useJpeg
// and encodes it to w.
func DrawFunction(w io.Writer, f Function, width, height int, viewport image.Rectangle, useJpeg bool) {
	m := image.NewRGBA(image.Rect(0, 0, width, height))

	drawBackground(m)
	drawAxes(m, viewport)
	drawValues(m, viewport, f)
	if useJpeg {
		jpeg.Encode(w, m, nil)
	} else {
		png.Encode(w, m)
	}
}
